"""图像的线性变换

目前方法仅支持操作灰度图，操作彩色图像需对图像的多个通道处理后进行手动混合

操作方法
1、采样
2、
"""

from PIL import Image

import numpy as np
import time

# 基础方法


def coordinate_match(origin_shape: tuple, output_shape: tuple) -> np.ndarray:
    """将新图像像素坐标与原图像像素坐标匹配

    将图像B的尺寸变换到与图像A大小一致，此时像素阵列B的最后一列和最后一行与像素阵列A的最后一列和最后一行对齐，但像素间距不相等。
    输出的map将一一表示图像B中每个像素在A坐标系中的坐标。
    此map将用于插值时灰度值的计算。

    Args:
        origin_shape: 原图像高m1、宽n1
        output_shape: 新图像高m2、宽n2

    Returns:
        shape为 (m2, n2, 2) 的阵列，里面的每个数对都是新图像对应点在原图像坐标系中的坐标
    """
    m1, n1 = origin_shape
    m2, n2 = output_shape
    output_map = np.ndarray((m2, n2, 2), dtype=np.double)

    # 输出图像四个顶点坐标
    # 左上、左下、右上、右下
    output_map[0, 0] = (0, 0)
    output_map[m2 - 1, 0] = (m1 - 1, 0)
    output_map[0, n2 - 1] = (0, n1 - 1)
    output_map[m2 - 1, n2 - 1] = (m1 - 1, n1 - 1)

    # 填充头尾两列坐标
    deltax = (m1 - 1) / (m2 - 1)
    for i in range(1, m2 - 1):
        output_map[i, 0] = (i * deltax, 0)
        output_map[i, n2 - 1] = (i * deltax, n1 - 1)

    # 填充中间部分坐标
    deltay = (n1 - 1) / (n2 - 1)
    for i in range(m2):
        for j in range(1, n2 - 1):
            output_map[i, j] = (output_map[i, 0, 0], j * deltay)

    return output_map


def linear_interpolation(x, x1, x2, f1, f2):
    """线性内插

    x在x1和x2之间。f1 = f(x1), f2 = f(x2)。
    """
    return (f2 - f1) / (x2 - x1) * (x - x1) + f1


# 图片采样类，上采样下采样均支持
class ImageSample(object):
    """图片采样方法集合"""

    @staticmethod
    def nearest_neighbor_interpolation(
        im: Image.Image, output_shape: tuple
    ) -> Image.Image:
        """最邻内插法。

        将原图像中位置最接近的像素的灰度值赋给新图像。
        优缺点：简单，但是可能导致某些直边缘的严重失真。

        Args:
            im: 输入图像
            output_shape: 输出图像的尺寸，格式 (m, n)

        Returns:
            插值之后的灰度图像
        """
        input_array = np.asarray(im.convert("L"), dtype=np.uint8)
        input_shape = input_array.shape
        output_array = np.zeros(output_shape, dtype=np.uint8)

        output_map = coordinate_match(input_shape, output_shape)

        for i in range(output_shape[0]):
            for j in range(output_shape[1]):
                # 使用np.round()计算很慢
                # x, y = output_map[i, j]
                # output_array[i, j] = input_array[int(np.round(x)), int(np.round(y))]

                # 自己判断四舍五入比上面的代码快很多
                x, y = output_map[i, j]
                x_nearest = int(x) if x - int(x) < 0.5 else int(x) + 1
                y_nearest = int(y) if y - int(y) < 0.5 else int(y) + 1
                output_array[i, j] = input_array[x_nearest, y_nearest]

        return Image.fromarray(output_array, mode="L")

    @staticmethod
    def bilinear_interpolation(im: Image.Image, output_shape: tuple) -> Image.Image:
        """双线性内插。

        实际上双线性内插并不属于线性变换，因为存在xy项。

        等同于先从x方向插值，再从y方向插值（或先从y方向再从x方向）。
        由以下公式求出，其中 v(x, y)表示插值点的灰度值，x、y代表该点位置。
        v(x, y) = ax + by + cxy + d
        """

        def _bilinear_interpolation_one_point(x, y, x1, y1, x2, y2, f1, f2, f3, f4):
            """

            插值所用的四个已知点的位置分别为左上、左下、右上、右下，
            坐标分别为(x1, y1)、(x2, y1)、(x1, y2)、(x2, y2)
            四个点灰度值分别为f1、f2、f3、f4
            插值点坐标为(x, y)

            Returns:
                插值点的灰度值
            """
            # delta_x = x2 - x1
            # delta_y = y2 - y1
            # delta_xy = delta_x * delta_y
            # delta_f42 = f4 - f2
            # delta_f31 = f3 - f1
            # c = (delta_f42 - delta_f31) / delta_xy
            # a = -y1 * (delta_f42 - delta_f31) / delta_xy
            # b = (-x1 * delta_f42 + x2 * delta_f31) / delta_xy
            # d = (
            #     -f2 * delta_y * x1 + f1 * delta_y * x1 - delta_x * y1 * delta_f31
            # ) / delta_xy + f1

            # v = a * x + b * y + c * x * y + d
            # if v > 255:
            #     print(f"v: {v}")

            # 方法2
            delta_x = x2 - x1
            delta_y = y2 - y1
            delta_f31 = f3 - f1
            delta_f42 = f4 - f2
            R1 = (y - y1) * delta_f31 / delta_y + f1
            R2 = (y - y1) * delta_f42 / delta_y + f2

            delta_R = R2 - R2
            v = (x - x1) * delta_R / delta_x + R1
            return int(v)

        input_array = np.asarray(im.convert("L"), dtype=np.int16)
        output_array = np.ndarray(output_shape, dtype=np.uint8)
        output_map = coordinate_match(input_array.shape, output_shape)

        # 计算左上部分像素灰度值（右边界和下边界需要特殊处理）
        for i in range(0, output_shape[0] - 1):
            for j in range(0, output_shape[1] - 1):
                x, y = output_map[i, j]
                x1, x2, y1, y2 = int(x), int(x) + 1, int(y), int(y) + 1
                f1 = input_array[x1, y1]
                f2 = input_array[x2, y1]
                f3 = input_array[x1, y2]
                f4 = input_array[x2, y2]
                output_array[i, j] = _bilinear_interpolation_one_point(
                    x, y, x1, y1, x2, y2, f1, f2, f3, f4
                )

        # 计算下边缘和右边缘像素灰度值，TODO：不好看
        for j in range(0, output_shape[1] - 1):
            x, y = output_map[-1, j]
            y1 = int(y)
            y2 = int(y) + 1
            f1 = input_array[-1, y1]
            f2 = input_array[-1, y2]
            output_array[-1, j] = int(linear_interpolation(y, y1, y2, f1, f2))

        for i in range(0, output_shape[0] - 1):
            x, y = output_map[i, -1]
            x1 = int(x)
            x2 = int(x) + 1
            f1 = input_array[x1, -1]
            f2 = input_array[x2, -1]
            output_array[i, -1] = int(linear_interpolation(x, x1, x2, f1, f2))

        output_array[-1, -1] = input_array[-1, -1]

        return Image.fromarray(output_array)

    pass


if __name__ == "__main__":
    with Image.open("~/Downloads/吔屎啦你.jpeg") as im:
        # im_sample1 = ImageSample.bilinear_interpolation(im, (1080, 1920))
        # im_sample1.show("1")
        start = time.time()
        im_sample2 = ImageSample.nearest_neighbor_interpolation(im, (1080, 1920))
        print(time.time() - start)
        im_sample2.show("2")
        im.convert("L").show("3")
